"""
Model Management Interface for Continual Learning

This interface defines methods for managing neural network models in continual learning,
including architecture updates, feature extraction, and model state management.
"""

from abc import ABC, abstractmethod
from typing import Any, Dict, Optional, Union
import torch
import torch.nn as nn


class ModelManager(ABC):
    """Interface for managing neural network models in continual learning"""

    @abstractmethod
    def update_fc(self, nb_classes: int, **kwargs) -> None:
        """
        Update the final classification layer for new classes

        Args:
            nb_classes: Number of new classes to add
            **kwargs: Additional parameters for FC layer configuration
        """
        pass

    @abstractmethod
    def extract_features(self, inputs: torch.Tensor) -> torch.Tensor:
        """
        Extract feature representations from inputs

        Args:
            inputs: Input tensor

        Returns:
            Feature representations
        """
        pass

    @abstractmethod
    def extract_token(self, inputs: torch.Tensor) -> torch.Tensor:
        """
        Extract token representations (for transformer-based models)

        Args:
            inputs: Input tensor

        Returns:
            Token representations
        """
        pass

    @abstractmethod
    def get_feature_dim(self) -> int:
        """Get the dimensionality of feature representations"""
        pass

    @abstractmethod
    def save_checkpoint(self, filepath: str) -> None:
        """
        Save model checkpoint

        Args:
            filepath: Path to save the checkpoint
        """
        pass

    @abstractmethod
    def load_checkpoint(self, filepath: str) -> None:
        """
        Load model checkpoint

        Args:
            filepath: Path to load the checkpoint from
        """
        pass

    @abstractmethod
    def freeze_feature_extractor(self) -> None:
        """Freeze the feature extraction layers"""
        pass

    @abstractmethod
    def unfreeze_feature_extractor(self) -> None:
        """Unfreeze the feature extraction layers"""
        pass


class NetworkStateManager(ABC):
    """Interface for managing network state and versioning"""

    @abstractmethod
    def backup_network(self) -> None:
        """Create a backup of the current network state"""
        pass

    @abstractmethod
    def restore_network(self, version: Optional[str] = None) -> None:
        """
        Restore network from backup

        Args:
            version: Specific version to restore (if None, restore latest)
        """
        pass

    @abstractmethod
    def get_network_state(self) -> Dict[str, Any]:
        """Get current network state dictionary"""
        pass

    @abstractmethod
    def set_network_state(self, state: Dict[str, Any]) -> None:
        """Set network state from dictionary"""
        pass


class ModelArchitectureAdapter(ABC):
    """Interface for adapting model architecture during continual learning"""

    @abstractmethod
    def expand_architecture(self, expansion_config: Dict[str, Any]) -> None:
        """
        Expand model architecture for new tasks

        Args:
            expansion_config: Configuration for architecture expansion
        """
        pass

    @abstractmethod
    def add_task_specific_modules(self, task_id: int, module_config: Dict[str, Any]) -> None:
        """
        Add task-specific modules to the model

        Args:
            task_id: Task identifier
            module_config: Configuration for task-specific modules
        """
        pass

    @abstractmethod
    def get_task_modules(self, task_id: int) -> Dict[str, nn.Module]:
        """Get modules specific to a task"""
        pass


class FeatureProcessor(ABC):
    """Interface for processing and transforming features"""

    @abstractmethod
    def normalize_features(self, features: torch.Tensor) -> torch.Tensor:
        """
        Normalize feature vectors

        Args:
            features: Input features

        Returns:
            Normalized features
        """
        pass

    @abstractmethod
    def compute_feature_statistics(self, features: torch.Tensor) -> Dict[str, torch.Tensor]:
        """
        Compute statistics for features

        Args:
            features: Input features

        Returns:
            Dictionary containing feature statistics
        """
        pass

    @abstractmethod
    def apply_feature_transform(
        self,
        features: torch.Tensor,
        transform_type: str
    ) -> torch.Tensor:
        """
        Apply transformation to features

        Args:
            features: Input features
            transform_type: Type of transformation to apply

        Returns:
            Transformed features
        """
        pass


class ModelEnsemble(ABC):
    """Interface for managing model ensembles"""

    @abstractmethod
    def add_model(self, model: nn.Module, weight: float = 1.0) -> None:
        """Add a model to the ensemble"""
        pass

    @abstractmethod
    def remove_model(self, model_id: str) -> None:
        """Remove a model from the ensemble"""
        pass

    @abstractmethod
    def ensemble_predict(self, inputs: torch.Tensor) -> torch.Tensor:
        """Make predictions using the ensemble"""
        pass

    @abstractmethod
    def update_weights(self, weights: Dict[str, float]) -> None:
        """Update ensemble weights"""
        pass
